import os
import cv2
import numpy as np
import matplotlib.pyplot as plt


def load_image(image_dir, file_name):
    image_path = os.path.join(image_dir, file_name)
    image = cv2.imread(image_path)
    return image


def visualize_images_with_opencv(center_image, sub_image, sub_image_region, combined_mask, title):
    cv2.imshow(f"Center Image - {title}", center_image)
    cv2.imshow(f"Sub Image - {title}", sub_image)
    cv2.imshow(f"Sub Image Region - {title}", sub_image_region.astype(np.uint8))
    cv2.imshow(f"Combined Mask - {title}", (combined_mask * 255).astype(np.uint8))
    cv2.waitKey(0)
    cv2.destroyAllWindows()


def merge_images(center_image, neighbor_images, offsets, visualize=False):
    merged_image = center_image.copy()
    center_region = center_image[512:1536, 512:1536]

    for idx, (image, (dx, dy)) in enumerate(zip(neighbor_images, offsets)):
        if image is not None:
            x_start, x_end = 512 + dx, 1536 + dx
            y_start, y_end = 512 + dy, 1536 + dy

            # Ensure valid coordinates within the image dimensions
            x_start_clipped = max(0, x_start)
            x_end_clipped = min(image.shape[1], x_end)
            y_start_clipped = max(0, y_start)
            y_end_clipped = min(image.shape[0], y_end)

            sub_image = image[y_start_clipped:y_end_clipped, x_start_clipped:x_end_clipped]

            # Calculate the corresponding position in the center region
            center_x_start = x_start_clipped - dx
            center_y_start = y_start_clipped - dy
            center_x_end = x_end_clipped - dx
            center_y_end = y_end_clipped - dy

            # Adjust to ensure within bounds
            center_x_start = max(512, center_x_start)
            center_y_start = max(512, center_y_start)
            center_x_end = min(1536, center_x_end)
            center_y_end = min(1536, center_y_end)

            sub_image_region = center_region[
                               (center_y_start - 512):(center_y_end - 512),
                               (center_x_start - 512):(center_x_end - 512)
                               ]

            mask_center = (sub_image_region != [0, 0, 0])
            mask_sub = (sub_image != [0, 0, 0])
            combined_mask = mask_center & mask_sub

            # 合并规则
            sub_image_region[~mask_center & mask_sub] = sub_image[~mask_center & mask_sub]
            sub_image_region[mask_center & ~mask_sub] = sub_image_region[mask_center & ~mask_sub]
            # sub_image_region[combined_mask] = ((sub_image_region[combined_mask].astype(np.float32) + sub_image[
            #     combined_mask].astype(np.float32)) / 2).astype(np.uint8)
            sub_image_region[combined_mask] = np.maximum(sub_image_region[combined_mask],
                                                         sub_image[combined_mask])

            if visualize:
                visualize_images_with_opencv(center_image, sub_image, sub_image_region, combined_mask,
                                             f"Step {idx + 1}")

    merged_image[512:1536, 512:1536] = center_region
    return merged_image


def get_neighbor_images(image_dir, row, col, prefix):
    neighbor_coords = [
        (row - 1, col - 1), (row - 1, col), (row - 1, col + 1),
        (row, col - 1), (row, col + 1),
        (row + 1, col - 1), (row + 1, col), (row + 1, col + 1)
    ]
    offsets = [
        (1024, 1024), (0, 1024), (-1024, 1024),
        (1024, 0), (-1024, 0),
        (1024, -1024), (0, -1024), (-1024, -1024)
    ]

    neighbor_images = []
    for nr, nc in neighbor_coords:
        neighbor_file = f"{prefix}tr{nr}-tc{nc}.png"
        neighbor_image = load_image(image_dir, neighbor_file) if os.path.exists(
            os.path.join(image_dir, neighbor_file)) else None
        neighbor_images.append(neighbor_image)

    return neighbor_images, offsets


def process_images(image_dir, output_dir, visualize=False):
    os.makedirs(output_dir, exist_ok=True)

    for image_file in os.listdir(image_dir):
        if image_file.endswith('.png'):
            parts = image_file.split('tr')[1].split('-tc')
            row = int(parts[0])
            col = int(parts[1].split('.')[0])
            prefix = image_file.split('tr')[0]

            # 加载中心图像
            center_image = load_image(image_dir, image_file)

            # 获取相邻图像和偏移量
            neighbor_images, offsets = get_neighbor_images(image_dir, row, col, prefix)

            # 合并图像
            merged_image = merge_images(center_image, neighbor_images, offsets, visualize=visualize)

            # 裁剪中间区域并保存
            final_image = merged_image[512:1536, 512:1536]
            save_tile(final_image, output_dir, image_file)


def save_tile(image, output_dir, file_name):
    output_path = os.path.join(output_dir, file_name)
    cv2.imwrite(output_path, image)
# 示例用法
image_dir = 'E:\\wafer52\\11867_32nm_coarse_align'  # 替换为实际扩展图像块目录
output_dir = 'E:\\wafer52\\11867_32nm_coarse_align_merge'  # 替换为实际输出目录

process_images(image_dir, output_dir, False)
